#include <iostream>
#include <fstream>
#include <malloc.h>
#include <stdio.h>
#include <stdlib.h>

using namespace std;

#include "types.h"
#include "rsakeys.h"
#include "base64.h"
#include "cmdargs.h"
#include "crypto.h"
#include "helper.h"

extern cmdargs args;

static char msg_error[] = "message too long";
static uchar decode_error[] = "decode error";

bool rsaVerify()
{
	cout << "verifying signature " << endl;
 	if (!rsaDecrypt())
	 return FALSE;

	system ("md5sum plaintext.txt > orig.md5");
	
	ifstream origmd5("orig.md5", ios :: binary);
	char  orig[32];
	origmd5.read((char *)orig, 32);
	origmd5.close();

	ifstream is (args.outfile.c_str(), ios::binary);
	uchar decrypt[32];
	is.read((char *)decrypt, 32);
	is.close();
	
	cout << "orig to md5: " << string ((char*)orig, 32)<< endl;
	cout << "md5 after decrypt: " <<string ((char *)decrypt, 32) << endl;

	if ( 0 == memcmp(orig, decrypt, 32))
		cout <<"matched found : VERIFIED" << endl;
	else
		cout <<"do not match: Signature cannot be verified."<< endl;

 return TRUE;
}



// Encrypt using public key and decrypt using private key.
bool rsaEncrypt()
{
	cout << "encrypting file" << endl;
	mpz_t modulus, e;
	mpz_init(modulus);
	mpz_init(e);
	
	createDerFile();
	getPublicKey(modulus, e);
	
	int octetsinmod = mpz_sizeinbase(modulus, 256);
	ifstream is (args.infile.c_str(), ios:: binary);
	char *msg;
	
	// read the complete private key file
	is.seekg(0, ios::end);
	uint len = is.tellg();
	is.seekg(0, ios::beg);
	msg = (char *)malloc(len); // buffer to read in the private key file.
	is.read(msg, len);
	is.close();

	char *em;
	em = eme_pkcs1_v1_5_encode(string(msg, len), octetsinmod-1);
	if (!strcmp(em, msg_error))
	{
		printf("%s\n",em );
		return FALSE;
	}

	// convert the EM to a number;
	mpz_t msginint, cipherinint;
	mpz_init (msginint);
	mpz_init (cipherinint);

	//mpz_import(msginint, octetsinmod, 1,1,0,0, (char *)em);
	mpz_import(msginint, octetsinmod -1, 1,1,0,0, (char *)em);
	mpz_powm(cipherinint, msginint, e, modulus);
	
	uint count;
	char cipher[1024];
	mpz_export((void *)cipher, &count, 1,1,0,0, cipherinint);
	
	ofstream os (args.outfile.c_str(), ios::binary);

	if (!os.is_open())
	{	
		cout << "unable to write to cipher file" << endl;
		return FALSE;
	}

	os.write(cipher, count);
	
	os.close();
	free (msg);
	mpz_clear (modulus);
	mpz_clear (e);
	mpz_clear (msginint);
	mpz_clear(cipherinint);
	return TRUE;
}


bool rsaDecrypt()
{
	cout << "decrypting file" << endl;
	mpz_t modulus, key;
	mpz_init(modulus);
	mpz_init(key);
	
	createDerFile();
	if (args.keyfromcert == TRUE)
		getPubKeyFromCert(modulus, key);
	else
		getPrivateKey(modulus, key);

	int octetsinmod = mpz_sizeinbase(modulus, 256);
	ifstream is (args.infile.c_str(), ios:: binary);
	char *cipher;
	
	// read the complete private key file
	is.seekg(0, ios::end);
	uint len = is.tellg();
	is.seekg(0, ios::beg);
	cipher = (char *)malloc(len); // buffer to read in the private key file.
	is.read(cipher, len);
	is.close();
	
	mpz_t cipherinint, plaininint;
	mpz_init(cipherinint);
	mpz_init(plaininint);

	mpz_import(cipherinint, octetsinmod , 1,1,0,0, cipher);
	mpz_powm(plaininint, cipherinint, key, modulus);
	
	uchar plain[512];
	uint count;
	mpz_export((void *)plain, &count, 1,1,0,0, plaininint);

/*
	ofstream t2 ("decryptedEM", ios :: binary);
	t2.write((char *)plain, count);
	t2.close();
*/
	uchar *msg;
	if (args.keyfromcert == TRUE)
	{
		msg = signed_msg_decode(plain, count);
	}
	else
	{
		count --;
		msg = eme_pkcs1_v1_5_decode(plain, count);

		if (!strcmp ((const char *)msg, (const char *)decode_error))
		{
			cout << msg << endl;
			return FALSE;
		}
	}

	// write messge to the file.
	ofstream os(args.outfile.c_str(), ios::binary);
	if (!os.is_open())
	{
		cout << "cannot open the output file" << endl;
		return FALSE;
	}
	os.write ((char *)msg, count);

	os.close();
	mpz_clear(modulus);
	mpz_clear(key);
	mpz_clear(plaininint);
	mpz_clear(cipherinint);	
	free (cipher);
	
	return TRUE;
}

uchar *signed_msg_decode(uchar *em, uint &count)
{
	
	uchar *buff = (uchar *)malloc(16);
	int i,j;
	for ( i = 0, j = count-16; j < count ; j++, i++)
		buff[i] = em[j];

/*
	ofstream test1("dump1", ios::binary);
	test1.write((char *)buff, 16);
	test1.close();
*/
	uchar *msg = (uchar *)malloc(32);
	for (i = 0, j = 0; i < 16; i++, j+=2)
	{
		sprintf ((char *)&msg[j], "%02x", buff[i]);
	}
	//cout << "msg: "<< string ((char *)msg, 32)<< endl;
	count = 32;
	free (buff);
	return msg;
}

uchar *eme_pkcs1_v1_5_decode(uchar *em, uint &count)
{
	uint emsize = count;

	if (emsize < 10)
		return 	decode_error;

	if (em[0] != 0x02)
		return decode_error;

	int j;
	for ( j = 0; (j < count) && (em[j]) ; j ++)
		;
		
	uchar *msg = (uchar *)malloc(128);
	memcpy(msg, &em[j+1], count -j - 1);

	count = count -j - 1;
	return msg;
}

char *eme_pkcs1_v1_5_encode(string msg, int emlen)
{
	uint msgsize = msg.length();

	if (msgsize > (emlen-10))
		return 	msg_error;

	uint nofpadbytes = emlen - msgsize - 2;
	char *em;
	em = (char *)malloc(emlen);

	em[0] = 0x02;
	string ps;
	ps = getRandomString(nofpadbytes);
	memcpy(&em[1], ps.c_str(), nofpadbytes);
	em[nofpadbytes+1] = '\0';
	memcpy(&em[nofpadbytes+2], msg.c_str(), msgsize );

	return em;
}

bool createDerFile()
{
	base64 b;

	// If input is not in DER, convert to DER. 
	if (args.inkey.find(".der") == string::npos )
	{
		string derfile = args.inkey;
		derfile.replace(derfile.end()-3, derfile.end(), "der");
		
		ifstream testder(derfile.c_str());
		if (!testder)
		{
			if (!b.decodeKeyFromBase64(args.inkey.c_str(), derfile.c_str()))
			{	
				cout << "unable to convert PEM to DER " << endl;
				return FALSE;
			}
		}
		else
		{
			testder.close();
		}
		args.inkey = derfile;
	}
	return TRUE;
}


// assumes args.inkey points to a DER Encoded file.
bool getPrivateKey(mpz_t modulus, mpz_t d)
{
	base64 b;

	ifstream is (args.inkey.c_str(), ios::binary);
	if (!is.is_open())
	{
		cout << "Error: cannot open input private key file: " << args.inkey << endl;
		return FALSE;
	} 
	uchar *inbuff; // to read from the private file

	// read the complete private key file
	is.seekg(0, ios::end);
	uint len = is.tellg();
	is.seekg(0, ios::beg);
	inbuff = (uchar *)malloc(len); // buffer to read in the private key file.
	is.read((char *)inbuff, len);
	is.close();
	
	uint index = 0;
	uint noflenoctets = 0;
	uint contentoctets = 0;
	uchar firstlenoctet = '\0';

	// skip the header
	index = 1;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}

	// skip the version info
	index += 3;
	
	// reached modulus.
	if (inbuff[index++] != 2)
	{
		cout<< "error in private file" << endl;
		return FALSE;
	}
	
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
			
		noflenoctets += 1;
	}
	else
	{
		contentoctets = firstlenoctet;
	}

	//index points to first content octet
	uchar *keybuff = (uchar *)malloc (contentoctets);
	memcpy (keybuff, &inbuff[index], contentoctets);

	mpz_import(modulus, contentoctets, 1,1,1,0, keybuff);

	// skip the public component.
	index += contentoctets;
	index += 5;
	
	// extract private key component.
	if (inbuff[index++] != IDENT_INTEGER)
	{
		cout << " error reading private key file" << endl;
		return FALSE;
	}
	
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
			
		noflenoctets += 1;
	}
	else
	{
		contentoctets = firstlenoctet;
	}

	//index points to first content octet
	memcpy (keybuff, &inbuff[index], contentoctets);

	mpz_import(d, contentoctets, 1,1,1,0, keybuff);
	free (inbuff);
	free (keybuff);

	return TRUE;
}


// assumes args.inkey points to a DER Encoded file.
bool getPublicKey(mpz_t modulus, mpz_t e)
{
	base64 b;

	ifstream is (args.inkey.c_str(), ios::binary);
	if (!is.is_open())
	{
		cout << "Error: cannot open input private key file: " << args.inkey << endl;
		return FALSE;
	}

	uchar *inbuff; // to read from the public key file

	// read the complete public key file
	is.seekg(0, ios::end);
	uint len = is.tellg();
	is.seekg(0, ios::beg);
	inbuff = (uchar *)malloc(len); // buffer to read in the private key file.
	is.read((char *)inbuff, len);
	is.close();
	
	uint index = 0;
	uint noflenoctets = 0;
	uint contentoctets = 0;
	uchar firstlenoctet = 0;

	// skip sequence1
	index = 1;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}

	// skip the public key object identifier.
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
		noflenoctets += 1; // to include the first length octet.
	}
	else
	{
		noflenoctets = 1;
		contentoctets = firstlenoctet;
	}
	index += contentoctets;	

	// skip bitstring header.
	index ++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}
	//	skip pad byte
	index ++;

	// skip sequence header
	index ++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}

	// reached modulus.
	if (inbuff[index++] != 2)
	{
		cout<< "error in public file" << endl;
		return FALSE;
	}
	
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
			
		noflenoctets += 1;
	}
	else
	{
		contentoctets = firstlenoctet;
	}

	//index points to first content octet
	uchar *keybuff = (uchar *)malloc (contentoctets);
	memcpy (keybuff, &inbuff[index], contentoctets);
	
	mpz_import(modulus, contentoctets, 1,1,0,0, keybuff);

	index += contentoctets;
	
	// extract public key component.
	if (inbuff[index++] != IDENT_INTEGER)
	{
		cout << " error reading public key file" << endl;
		return FALSE;
	}
	
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
			
		noflenoctets += 1;
	}
	else
	{
		contentoctets = firstlenoctet;
	}

	//index points to first content octet
//	char * newbuff[1024];
	memcpy (keybuff, &inbuff[index], contentoctets);

	mpz_import(e, contentoctets, 1,1,0,0, keybuff);
	free (inbuff);
	free (keybuff);

	return TRUE;
}


bool getPubKeyFromCert(mpz_t modulus, mpz_t e)
{
	base64 b;

	ifstream is (args.inkey.c_str(), ios::binary);
	if (!is.is_open())
	{
		cout << "Error: cannot open input private key file: " << args.inkey << endl;
		return FALSE;
	}

	uchar *inbuff; // to read from the public key file

	// read the complete public key file
	is.seekg(0, ios::end);
	uint len = is.tellg();
	is.seekg(0, ios::beg);
	inbuff = (uchar *)malloc(len); // buffer to read in the private key file.
	is.read((char *)inbuff, len);
	is.close();
	
	/*
	ofstream t1("certinder", ios:: binary);
	t1.write((char *)inbuff, len);
	t1.close();
*/

	uint index = 0;
	uint noflenoctets = 0;
	uint contentoctets = 0;
	uchar firstlenoctet = 0;

	// enter  sequence1
	if (inbuff[index++] != IDENT_SEQ)
	{
		cout << "error in seq1 header in cert" << endl;
		return FALSE;
	}
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}

	// enter sequence2
	if (inbuff[index++] != IDENT_SEQ)
	{
		cout << "error in seq2 header in cert" << endl;
		return FALSE;
	}
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}

	// skip integer
	if (inbuff[index++] != IDENT_INTEGER)
	{
		cout << "error in integer header in cert" << endl;
		return FALSE;
	}
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
	}
	else
		contentoctets = firstlenoctet;

		index += contentoctets;

	// skip sequence3
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
	}
	else
		contentoctets = firstlenoctet;

		index += contentoctets;

	// skip sequence4 (oids)
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
	}
	else
		contentoctets = firstlenoctet;

		index += contentoctets;

	//skip sequence5 (time)
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
	}
	else
		contentoctets = firstlenoctet;

		index += contentoctets;

	// skip sequence6 
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
	}
	else
		contentoctets = firstlenoctet;

		index += contentoctets;

	// enter sequence 7. this has public key and mod.
	// skip the length octets.
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}

	// skip oid sequence
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
	}
	else
		contentoctets = firstlenoctet;

		index += contentoctets;

	// skip bitstring header	
	index ++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}
	//	skip pad byte
	index ++;

	// enter the sequence.
	index++;
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}
 
 // reached modulus.
	if (inbuff[index++] != 2)
	{
		cout<< "error reading modulus from certificate" << endl;
		return FALSE;
	}
	
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
			
		noflenoctets += 1;
	}
	else
	{
		contentoctets = firstlenoctet;
	}

	//index points to first content octet
	char *keybuff = (char *)malloc (contentoctets);
	memcpy (keybuff, &inbuff[index], contentoctets);
	
	mpz_import(modulus, contentoctets, 1,1,0,0, keybuff);

	index += contentoctets;
/*
	ofstream test1("dump1", ios::binary);
	test1.write(keybuff, contentoctets );
	test1.close();
	*/
	// extract public key component.
	if (inbuff[index++] != IDENT_INTEGER)
	{
		cout << "error reading public key component from cert" << endl;
		return FALSE;
	}
	
	firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index++];
		else if (noflenoctets == 2)
		{
			contentoctets = inbuff[index++];
			contentoctets = contentoctets*256 + inbuff[index++];
		}
			
		noflenoctets += 1;
	}
	else
	{
		contentoctets = firstlenoctet;
	}

	//index points to first content octet
	memcpy (keybuff, &inbuff[index], contentoctets);
	mpz_import(e, contentoctets, 1,1,0,0, keybuff);

/*
	ofstream test2("pubKeyinOctets", ios::binary);
	test2.write(keybuff, contentoctets );
	test2.close();
	*/
	
	free (inbuff);
	free (keybuff);

	return TRUE;
}


